package cn.zhuwj.autodiscovery;

import cn.zhuwj.autodiscovery.config.StartConfig;
import cn.zhuwj.autodiscovery.utils.Assert;
import cn.zhuwj.autodiscovery.utils.NetUtils;
import cn.zhuwj.autodiscovery.utils.StringUtils;
import cn.zhuwj.autodiscovery.core.TcpPortThread;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Arrays;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zhuwj.cn
 */
public class AutoDiscovery {

    private final Logger logger = LoggerFactory.getLogger(AutoDiscovery.class);

    private ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(30, 50,
            0L,TimeUnit.MILLISECONDS,
            new LinkedBlockingQueue<Runnable>());

    private StartConfig config;
    private StartConfig.Callback callback;
    private final AtomicInteger currIndex = new AtomicInteger(0);

    private AutoDiscovery(){
        this((StartConfig.Callback) null);
    }

    public AutoDiscovery(StartConfig config) {
        checkConfig(config);
        this.config = setDefaultValue(config);
        this.callback = this.config.getCallback();
    }

    public AutoDiscovery(StartConfig.Callback callback) {
        this.config = setDefaultValue(null);
        this.config.setCallback(callback);
        this.callback = this.config.getCallback();
    }

    public void start() throws InterruptedException {
        if(StringUtils.isNotBlank(config.getIp())){
            String [] fragments = config.getIp().split("\\.");
            int [] ips = new int[5];
            for(int i = 0;i < fragments.length;i++){
                ips[i] = Integer.parseInt(fragments[i]);
            }
            ips[4] = ips[3];
            TcpPortThread.IP_LIST = ips;
        }else{
            String [] fragments = config.getStartIp().split("\\.");
            int [] ips = new int[5];
            for(int i = 0;i < fragments.length;i++){
                ips[i] = Integer.parseInt(fragments[i]);
            }
            fragments = config.getEndIp().split("\\.");
            ips[4] = Integer.parseInt(fragments[3]);
            TcpPortThread.IP_LIST = ips;
        }
        Integer port = config.getPort();
        if(port != null){
            TcpPortThread.PORT_LIST = new int[] {port,port};
        }else{
            TcpPortThread.PORT_LIST = new int[] {config.getStartPort(),config.getEndPort()};
        }
        TcpPortThread.max_thread_count = config.getThreadCount();
        logger.info("目标主机地址范围：{}", Arrays.toString(TcpPortThread.IP_LIST));
        logger.info("目标主机端口范围：{}", Arrays.toString(TcpPortThread.PORT_LIST));
        logger.info("目标主机最大扫描线程数：{}", TcpPortThread.max_thread_count);
        int total_ip_count = TcpPortThread.IP_LIST[4] - TcpPortThread.IP_LIST[3] + 1;
        int total_port_count = TcpPortThread.PORT_LIST[1] - TcpPortThread.PORT_LIST[0] + 1;
        int total = total_ip_count * total_port_count;
        long startTime = System.currentTimeMillis();
        CountDownLatch latch = new CountDownLatch(config.getThreadCount());
        for (int i = 0; i < config.getThreadCount(); i++){
            threadPoolExecutor.execute(new TcpPortThread("ScanPortThread-" + i, i, total , currIndex, startTime,  latch, this.threadPoolExecutor, this.callback));
        }
        if(callback != null){
            callback.onStart(total);
        }
        logger.info("目标主机开始进行扫描~");
        latch.await();
        if(callback != null){
            threadPoolExecutor.shutdownNow();
            callback.onFinished(System.currentTimeMillis() - startTime);
        }
    }

    private StartConfig setDefaultValue(StartConfig config) {
        if(config == null){
            config = new StartConfig();
        }
        Integer threadCount = config.getThreadCount();
        if(threadCount == null || threadCount < 1 || threadCount > 50){
            config.setThreadCount(30);
        }
        if(StringUtils.isBlank(config.getIp())){
            if(StringUtils.isBlank(config.getStartIp())){
                String localIp = NetUtils.getRealIpv4IP();
                String startIp = localIp.replace(localIp.substring(localIp.lastIndexOf(".")),".1");
                config.setStartIp(startIp);
            }
            if(StringUtils.isBlank(config.getEndIp())){
                String localIp = NetUtils.getRealIpv4IP();
                String startIp = localIp.replace(localIp.substring(localIp.lastIndexOf(".")),".255");
                config.setEndIp(startIp);
            }
        }
        if(config.getPort() == null){
            if(config.getStartPort() == null){
                config.setStartPort(1);
            }
            if(config.getEndPort() == null){
                config.setEndPort(65535);
            }
        }
        return config;
    }

    private void checkConfig(StartConfig config) {
        if(config == null){
            return;
        }
        if(StringUtils.isNotBlank(config.getIp())){
            // 固定IP地址
            checkIpAddress(config.getIp());
        }else{
            // 范围IP地址
            if(StringUtils.isNotBlank(config.getStartIp())){
                checkIpAddress(config.getStartIp());
            }
            if(StringUtils.isNotBlank(config.getEndIp())){
                checkIpAddress(config.getEndIp());
            }
        }

        if(config.getPort() != null){
            // 固定端口
            checkPort(config.getPort());
        }else{
            // 范围端口
            if(config.getStartPort() != null){
                checkPort(config.getStartPort());
            }
            if(config.getEndPort() != null){
                checkPort(config.getEndPort());
            }
        }
    }

    private void checkIpAddress(String address){
        String [] ipFragments = address.split("\\.");
        Assert.isTrue(ipFragments.length == 4, "IP地址格式不正确 => ["+address+"]");
        for(String fragment : ipFragments){
            Assert.isTrue(StringUtils.isNumeric(fragment), "IP地址格式不正确 => ["+fragment+"]");
            int f = Integer.parseInt(fragment);
            Assert.isTrue(f >= 0 && f <= 255, "IP地址格式不正确 => ["+fragment+"]");
        }
    }

    private void checkPort(Integer port){
        Assert.isTrue(port >= 1, "端口范围不正确（1~65535） => ["+port+"]");
        Assert.isTrue(port <= 65535, "端口范围不正确（1~65535） => ["+port+"]");
    }

    public StartConfig getConfig() {
        return config;
    }

    public void setConfig(StartConfig config) {
        this.config = config;
    }

    public StartConfig.Callback getCallback() {
        return callback;
    }

    public void setCallback(StartConfig.Callback callback) {
        this.callback = callback;
        this.config.setCallback(callback);
    }

    public void stopScan(){
        threadPoolExecutor.shutdownNow();
    }

}
